package com.agentframework.core.environment;

import com.agentframework.core.agent.AbstractAgent;
import com.agentframework.core.agent.Agent;
import com.agentframework.core.agent.AgentHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.event.EventListener;
import org.springframework.lang.NonNull;

import java.util.*;
import java.util.concurrent.CountDownLatch;

import static com.agentframework.core.agent.Agent.*;

/**
 * @author wangjun
 **/
public abstract class AbstractEnvironment implements ApplicationContextAware, InitializingBean {

    protected static final Logger logger = LoggerFactory.getLogger(AbstractEnvironment.class);

    private final String envId;

    protected List<Agent> agents = new ArrayList<>();

    private final EnvironmentListeners listeners = new EnvironmentListeners();

    private final CountDownLatch completionLatch;

    private ApplicationContext applicationContext;

    /**
     * 环境结束条件列表
     */
    private final List<EnvironmentEndCondition> endConditions = new ArrayList<>();

    public AbstractEnvironment() {
        this.envId = "env_".concat(UUID.randomUUID().toString());
        this.completionLatch = new CountDownLatch(1);
    }

    public void publishMessage(Message message) {
        if (agents.isEmpty()) {
            logger.debug("env agents is empty");
            return;
        }

        if (message == null) {
            logger.error("env publish message is null");
            return;
        }

        // 如果满足环境结束条件 不再往agent队列发消息
        boolean anyConditionMet = endConditions.stream().anyMatch(c -> c.shouldEnd(this));
        if (anyConditionMet) {
            logger.info("Environment {} 已经满足关闭条件，不再接收消息", this.envId);
            return;
        }

        if (AgentHolder.getAgent() != null) {
            message.getMetadata().put(MSG_FROM, AgentHolder.getAgent().getAgentId());
        }

        applicationContext.publishEvent(new MessageEvent(message));
    }

    @EventListener
    public void onMessage(MessageEvent event) {
        if (event.getSource() instanceof Message message) {
            processMessage(message);
        }
    }

    protected void processMessage(Message message) {
        var sendToObj = message.getMetadata().get(SEND_TO);
        var sendTo = (sendToObj instanceof Set<?>) ? (Set<?>) sendToObj : new HashSet<>();
        if (sendTo.isEmpty()) {
            logger.debug("Message sendTo is not set, dont put into message queue, message {}", message);
            return;
        }

        if (sendTo.contains(TO_ALL)) {
            getAgents().forEach(agent -> deliverMessage(agent, message));
        } else {
            getAgents().stream().filter(agent -> sendTo.contains(agent.getAgentId())).forEach(agent -> deliverMessage(agent, message));
        }
    }

    protected void deliverMessage(Agent agent, Message message) {
        logger.debug(">>>>>put message to agent {} - {} process message {}", ((AbstractAgent) agent).getAgentName(), agent.getAgentId(), message);
        agent.getContext().putMessage(message);
    }

    public void run() {

        try {
            // 初始化
            init();

            // 启动
            start();

            // 开始异步任务和空闲检查
//            EnvironmentCompletionChecker completionChecker = new EnvironmentCompletionChecker(this);
//            completionChecker.startChecking();

            completionLatch.await();

        } catch (Exception e) {
            logger.error("Environment execution failed", e);
            // 错误事件
            error(e);
        } finally {
            // 关闭事件
            shutdown();
        }

    }

    protected void init() {
        listeners.init(this);
    }

    protected void start() {
        getAgents().forEach(agent -> agent.run(null));
        listeners.started(this);
    }

    protected void error(Exception e) {
        listeners.error(this, e);
    }

    protected void shutdown() {
        listeners.shutdown(this);
        completeEnvironment();
    }

    public void completeEnvironment() {
        completionLatch.countDown();
    }

    public String getEnvId() {
        return this.envId;
    }

    public void addAgent(Agent agent) {
        addAgents(List.of(agent));
    }

    public void addAgents(List<Agent> agents) {
        agents.forEach(agent -> agent.getContext().setEnvironment(this));
        this.agents.addAll(agents);
    }

    public List<Agent> getAgents() {
        return agents;
    }

    public void setAgents(List<Agent> agents) {
        agents.forEach(agent -> agent.getContext().setEnvironment(this));
        this.agents = agents;
    }

    public void addListener(EnvironmentEventListener<?> listener) {
        this.addListeners(List.of(listener));
    }

    public void addListeners(List<EnvironmentEventListener<?>> listeners) {
        this.listeners.addListeners(listeners);
    }

    public void addEndCondition(EnvironmentEndCondition condition) {
        this.endConditions.add(condition);
    }

    public List<EnvironmentEndCondition> getEndConditions() {
        return this.endConditions;
    }

    @Override
    public void afterPropertiesSet() throws Exception {

    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

}
